
#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Author: Combofish
# Filename: main.py

from icecream import ic
from torchvision import datasets, transforms
from tqdm import tqdm
import os


train_data = datasets.MNIST(root="./data/", train=True, download=False)
test_data = datasets.MNIST(root="./data/", train=False, download=False)

'''
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean = [0.5],std = [0.5])])
train_data = datasets.MNIST('./data/', train = True, transform = transform,download = False)
test_data = datasets.MNIST('./data/', train = False, transform = transform)
'''

saveDirTrain = './DataImages-Train'
saveDirTest = './DataImages-Test'

if not os.path.exists(saveDirTrain):
    os.mkdir(saveDirTrain)
if not os.path.exists(saveDirTest):
    os.mkdir(saveDirTest)

ic(len(train_data), len(test_data))
ic(train_data[0])
ic(train_data[0][0])


def save_img(data, save_path):
    for i in tqdm(range(len(data))):
        img, label = data[i]
        img.save(os.path.join(save_path, str(i) + '-label-' + str(label) + '.png'))


save_img(train_data, saveDirTrain)
save_img(test_data, saveDirTest)
